#include "nb_Server.h"
#include <assert.h>

NB_Server::NB_Server(int nPort)
{
	m_nPort = nPort;
	m_hCompletionPort = NULL;
	m_sListen = NULL;
	m_nServerThreadsNum = 0;
	m_bShutdown = false;
	m_nMaxConnections = 100;
	m_nRecvThreads = 2;
	m_nMaxHandleThreads = 4;

	m_pMsgManager = NULL;
	m_pRequest = NULL;
	m_pResponse = NULL;

	memset(&m_CallBack, NULL, sizeof(CALL_BACK));

	WSADATA wsaData;
	WSAStartup(MAKEWORD(2, 2), &wsaData);
}

NB_Server::~NB_Server()
{
	m_bShutdown = true;

	for(UINT i = 0; i < m_nServerThreadsNum; i++)
		PostQueuedCompletionStatus(m_hCompletionPort, 0, (ULONG_PTR)NULL, NULL);

	if(m_hCompletionPort) 
		CloseHandle(m_hCompletionPort);

	if(m_sListen)
		closesocket(m_sListen);

	if(m_pRequest) delete m_pRequest;
	m_pRequest = NULL;

	if(m_pResponse) delete m_pResponse;
	m_pResponse = NULL;

	if(m_pMsgManager) delete m_pMsgManager;
	m_pMsgManager = NULL;

	WSACleanup();
}

void NB_Server::SetMaxConnections(int nMaxConnections)
{
	m_nMaxConnections = nMaxConnections > 0 ? nMaxConnections : m_nMaxConnections;
}

void NB_Server::SetRecvThreads( int nThreads )
{
	if( nThreads > 0)
		m_nRecvThreads = nThreads;
}

void NB_Server::SetMaxHandleThreads( int nMaxThreads )
{
	if( nMaxThreads > 0)
		m_nMaxHandleThreads = nMaxThreads;
}

void NB_Server::SetSessionStart(lpfnSession lpfn)
{
	m_CallBack.SessionStart = lpfn;
}

void NB_Server::SetSessionHandle(lpfnSession lpfn)
{
	m_CallBack.SessionHandle = lpfn;
}

void NB_Server::SetSessionClose(lpfnSession lpfn)
{
	m_CallBack.SessionClose = lpfn;
}

void NB_Server::Run()
{
	m_pMsgManager = new NB_MsgManager(m_nMaxConnections);
	m_pRequest = new NB_Request();
	m_pResponse = new NB_Response();

	m_pResponse->m_lpMsgManager = m_pMsgManager;
	m_pHandle = new NB_Handle(m_nMaxHandleThreads);
	m_pHandle->m_lpRequest = m_pRequest;
	m_pHandle->m_lpResponse = m_pResponse;
	m_pHandle->m_DispatchFunc = m_CallBack.SessionHandle;

	m_sListen = ::WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
	if (m_sListen == INVALID_SOCKET)
	{
		printf("WSASocket() failed, error:%d\n", WSAGetLastError());
		return;
	}

	SOCKADDR_IN  local;
	local.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
	local.sin_family = AF_INET;
	local.sin_port = htons(m_nPort);
	if ( SOCKET_ERROR == bind(m_sListen, (struct sockaddr *)&local, sizeof(SOCKADDR_IN)) )
	{
		printf("bind() failed, error:%d\n", WSAGetLastError());
		return;
	}

	if ( SOCKET_ERROR == listen(m_sListen, SOMAXCONN) )
	{
		printf("listen() failed, error:%d\n", WSAGetLastError());
		return;
	}

	m_pResponse->m_hComPort = m_hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, m_nRecvThreads);
	if(!m_hCompletionPort)
	{
		printf("CreateIoCompletionPort() failed, error:%d\n", WSAGetLastError());
		return;
	}

	if( ! m_pHandle->Start() )
	{
		printf("Handle ThreadPool Start Failed, error:%d \n", WSAGetLastError());
		return;
	}
	
	for (int i = 0; i < m_nRecvThreads; i++)
	{
		if(QueueUserWorkItem(ServerThread, this, WT_EXECUTELONGFUNCTION))
			m_nServerThreadsNum ++;
	}

	DWORD	nThreadID;
	HANDLE	m_hAcceptThread;
	m_hAcceptThread = ::CreateThread(NULL, 0, AcceptThread, this, 0, &nThreadID);
	if(m_hAcceptThread == NULL)
	{
		printf("create accept thread failed\n");
		return;
	}
}

DWORD WINAPI NB_Server::AcceptThread(LPVOID lParam)
{
	NB_Server *	pThis = (NB_Server*)lParam;
	SOCKET			sClient;
	sockaddr_in		addrClient;
	int				addrLen = sizeof(addrClient);

	while( ! pThis->m_bShutdown )
	{

		sClient = WSAAccept(pThis->m_sListen, (struct sockaddr *)&addrClient, &addrLen, NULL, NULL);
		if (sClient == INVALID_SOCKET) {
			continue;
		}

		NB_Msg * lpMsg = pThis->m_pMsgManager->Create();
		if(lpMsg == NULL) {
			closesocket(sClient);
			continue;
		}

		lpMsg->socket = sClient;
		lpMsg->addr = addrClient;
		CreateIoCompletionPort((HANDLE)lpMsg->socket, pThis->m_hCompletionPort, (ULONG_PTR)lpMsg, 0);

		pThis->m_pRequest->m_lpMsg = lpMsg;
		pThis->m_pResponse->m_lpMsg = lpMsg;
		if(pThis->m_CallBack.SessionStart != NULL)
			pThis->m_CallBack.SessionStart(pThis->m_pRequest, pThis->m_pResponse);

		lpMsg->operation = OP_RECV;

		WSARecv(lpMsg->socket,
				&lpMsg->wsabuf,
				1,
				&lpMsg->dwBytesOfRecvd,
				&lpMsg->dwFlags,
				&lpMsg->overlapped,
				NULL);
	}

	return 0;
}

DWORD WINAPI NB_Server::ServerThread(LPVOID lParam)
{
	NB_Server *	pThis = (NB_Server *)lParam;
	NB_Msg *	lpMsg = NULL;
	OVERLAPPED	overlapped;
	DWORD		dwBytesTransferred;

	while( TRUE )
	{
		GetQueuedCompletionStatus(pThis->m_hCompletionPort,
								&dwBytesTransferred,
								(PULONG_PTR)&lpMsg,
								(LPOVERLAPPED *)&overlapped,
								INFINITE);

		if( lpMsg == NULL )
		{
			break;
		}

		switch(lpMsg->operation)
		{
		case OP_RECV:
			{
				pThis->m_pRequest->m_lpMsg = lpMsg;
				pThis->m_pResponse->m_lpMsg = lpMsg;

				if(dwBytesTransferred <= 0)
				{
					pThis->m_CallBack.SessionClose(pThis->m_pRequest, pThis->m_pResponse);

					closesocket(lpMsg->socket);
					pThis->m_pMsgManager->Free(lpMsg);
				}
				else
				{
					if(lpMsg->buffer[0] != '\0')
					{
						pThis->m_pHandle->SetInMsg(lpMsg);

		/*				if(pThis->m_CallBack.SessionHandle != NULL)*/
							//pThis->m_CallBack.SessionHandle(pThis->m_pRequest, pThis->m_pResponse);
					}

					if(lpMsg != NULL)
					{
						lpMsg->operation = OP_RECV;

						WSARecv(lpMsg->socket,
								&lpMsg->wsabuf,
								1,
								&lpMsg->dwBytesOfRecvd,
								&lpMsg->dwFlags,
								&lpMsg->overlapped,
								NULL);
					}
				}
			}
			break;
		}
	}

	return 0;
}